# encoding: UTF-8

from __future__ import division
import json
import shelve
import os
import traceback
from collections import OrderedDict
from queue import Queue, Empty
from threading import Thread
from pymongo.errors import DuplicateKeyError
from datetime import datetime, timedelta

from vnpy.event import Event
from vnpy.trader.vtEvent import (EVENT_TICK, EVENT_TRADE, EVENT_CONTRACT,
                                 EVENT_ORDER, EVENT_TIMER)
from vnpy.trader.vtFunction import getTempPath, getJsonPath, todayDate
from vnpy.trader.vtObject import (VtLogData, VtSubscribeReq, VtBarData, VtTickData,
                                  VtOrderReq, VtCancelOrderReq)
from vnpy.trader.vtConstant import (PRODUCT_OPTION, OPTION_CALL, OPTION_PUT,
                                    DIRECTION_LONG, DIRECTION_SHORT,
                                    OFFSET_OPEN, OFFSET_CLOSE,
                                    PRICETYPE_LIMITPRICE)
from vnpy.pricing import black, bs, crr, bsCython, crrCython
from vnpy.trader.vtUtility import BarGenerator

from .omBase import (OmOption, OmUnderlying, OmChain, OmPortfolio, OmVixCalculator,
                     EVENT_OM_LOG, EVENT_OM_STRATEGY, EVENT_OM_STRATEGYLOG, EVENT_OM_VIX,
                     TICK_DB_NAME, MINUTE_DB_NAME, OM_DB_NAME)
from .strategy import STRATEGY_CLASS

# 定价模型字典
MODEL_DICT = {}
MODEL_DICT['black'] = black
MODEL_DICT['bs'] = bs
MODEL_DICT['crr'] = crr
MODEL_DICT['bsCython'] = bsCython
MODEL_DICT['crrCython'] = crrCython


########################################################################
class OmEngine(object):
    """期权主引擎"""
    impvFileName = 'PricingImpv.vt'
    impvFilePath = getTempPath(impvFileName)

    # ----------------------------------------------------------------------
    def __init__(self, mainEngine, eventEngine):
        """Constructor"""
        self.mainEngine = mainEngine
        self.eventEngine = eventEngine

        self.portfolio = None
        self.optionContractDict = {}  # symbol:contract

        self.strategyEngine = OmStrategyEngine(self, eventEngine)

        self.vixEngine = None

        self.registerEvent()

        print('Custom OmEngine is running..')

    # ----------------------------------------------------------------------
    def registerEvent(self):
        """注册事件监听"""
        self.eventEngine.register(EVENT_CONTRACT, self.processContractEvent)

    # ----------------------------------------------------------------------
    def processTickEvent(self, event):
        """行情事件"""
        tick = event.dict_['data']
        self.portfolio.newTick(tick)

    # ----------------------------------------------------------------------
    def processTradeEvent(self, event):
        """成交事件"""
        trade = event.dict_['data']
        self.portfolio.newTrade(trade)

    # ----------------------------------------------------------------------
    def processContractEvent(self, event):
        """合约事件"""
        contract = event.dict_['data']
        if contract.symbol and contract.productClass == PRODUCT_OPTION:
            # print(contract.name, contract.strikePrice, contract.underlyingSymbol)
            self.optionContractDict[contract.symbol] = contract

    # ----------------------------------------------------------------------
    def subscribeEvent(self, symbol):
        """订阅对应合约的事件"""
        contract = self.mainEngine.getContract(symbol)
        if not contract:
            self.writeLog(u'行情订阅失败，找不到合约：%s' % symbol)
            return

        vtSymbol = contract.vtSymbol

        # 订阅行情
        req = VtSubscribeReq()
        req.symbol = contract.symbol
        req.exchange = contract.exchange
        self.mainEngine.subscribe(req, contract.gatewayName)

        # 订阅事件
        self.eventEngine.register(EVENT_TICK + vtSymbol, self.processTickEvent)
        self.eventEngine.register(EVENT_TRADE + vtSymbol, self.processTradeEvent)

    # ----------------------------------------------------------------------
    def initEngine(self, fileName):
        """初始化引擎"""
        if self.portfolio:
            return False

        f = open(fileName)
        setting = json.load(f)
        f.close()

        # 读取定价模型
        model = MODEL_DICT.get(setting['model'], None)
        if not model:
            self.writeLog(u'找不到定价模型%s' % setting['model'])
            return

        # 创建标的对象
        underlyingDict = OrderedDict()

        for underlyingSymbol in setting['underlying']:
            contract = self.mainEngine.getContract(underlyingSymbol)
            if not contract:
                self.writeLog(u'找不到标的物合约%s' % underlyingSymbol)
                continue

            detail = self.mainEngine.getPositionDetail(contract.vtSymbol)

            underlying = OmUnderlying(contract, detail)
            underlyingDict[underlyingSymbol] = underlying

        # 创建期权链对象并初始化
        chainList = []

        for d in setting['chain']:
            chainSymbol = d['chainSymbol']
            underlyingSymbol = d['underlyingSymbol']
            r = d['r']

            # 锁定标的对象
            underlying = underlyingDict.get(d['underlyingSymbol'], None)
            if not underlying:
                self.writeLog(u'%s期权链的标的合约%s尚未创建，请检查配置文件' % (chainSymbol, underlyingSymbol))
                continue

            # 创建期权对象并初始化
            callDict = {}
            putDict = {}

            # print(chainSymbol)

            for symbol, contract in self.optionContractDict.items():
                # print(symbol, contract.underlyingSymbol)
                # print(d['chainSymbol'])
                if contract.underlyingSymbol == d['chainSymbol']:
                    detail = self.mainEngine.getPositionDetail(contract.vtSymbol)
                    option = OmOption(contract, detail, underlying, model, r)

                    key = str(option.k)
                    if contract.optionType is OPTION_CALL:
                        if key in callDict:
                            key += 'a'
                        callDict[key] = option
                    else:
                        if key in putDict:
                            key += 'a'
                        putDict[key] = option

            # print(callDict.keys())
            # print(putDict.keys())

            # 期权排序
            strikeList = callDict.keys()
            strikeList.sort()
            # print(strikeList)
            callList = [callDict[k] for k in strikeList]
            putList = [putDict[k] for k in strikeList]

            # 创建期权链
            chain = OmChain(chainSymbol, callList, putList)
            chainList.append(chain)

            # 添加标的映射关系
            underlying.addChain(chain)

        # 创建持仓组合对象并初始化
        self.portfolio = OmPortfolio(setting['name'], model, underlyingDict.values(), chainList)

        # 载入波动率配置
        self.loadImpvSetting()

        # 订阅行情和事件
        for underlying in underlyingDict.values():
            self.subscribeEvent(underlying.vtSymbol)

        for chain in chainList:
            for option in chain.optionDict.values():
                self.subscribeEvent(option.vtSymbol)

        # 订阅合约后，开始计算波动率指数
        self.vixEngine = OmVixEngine(self, self.eventEngine)

        # 载入成功返回
        return True

    # ----------------------------------------------------------------------
    def loadImpvSetting(self):
        """载入波动率配置"""
        f = shelve.open(self.impvFilePath)

        for chain in self.portfolio.chainDict.values():
            for option in chain.optionDict.values():
                option.pricingImpv = f.get(option.symbol, 0)

        f.close()

    # ----------------------------------------------------------------------
    def saveImpvSetting(self):
        """保存波动率配置"""
        if not self.portfolio:
            return

        f = shelve.open(self.impvFilePath)

        for chain in self.portfolio.chainDict.values():
            for option in chain.optionDict.values():
                f[option.symbol] = option.pricingImpv

        f.close()

    # ----------------------------------------------------------------------
    def stop(self):
        """关闭函数"""
        self.saveImpvSetting()
        if self.vixEngine:
            self.vixEngine.stop()

    # ----------------------------------------------------------------------
    def writeLog(self, content):
        """发出日志 """
        log = VtLogData()
        log.logContent = content

        event = Event(EVENT_OM_LOG)
        event.dict_['data'] = log
        self.eventEngine.put(event)

    # ----------------------------------------------------------------------
    def adjustR(self):
        """调整折现率"""
        if self.portfolio:
            self.portfolio.adjustR()

        for chain in self.portfolio.chainDict.values():
            self.writeLog(u'期权链%s的折现率r拟合为%.3f' % (chain.symbol, chain.r))


class OmVixEngine(object):
    """波动率计算引擎"""

    def __init__(self, omEngine, eventEngine):
        self.omEngine = omEngine
        self.mainEngine = omEngine.mainEngine
        self.strategyEngine = omEngine.strategyEngine
        self.eventEngine = eventEngine
        self.portfolio = omEngine.portfolio

        self.bg = BarGenerator(self.onVixBar)

        self.active = False
        self.queue = Queue()
        self.thread = Thread(target=self.run)

        self.vixDict = OrderedDict()
        self.timerDict = dict()

        for chain in self.portfolio.chainDict.values():
            self.vixDict[chain.symbol] = OmVixCalculator(chain)

        self.registerEvent()
        self.start()

    def start(self):
        """启动"""
        self.active = True
        self.thread.start()

    def stop(self):
        """关闭"""
        if self.active:
            self.active = False
            self.thread.join()

    def run(self):
        """运行入库的线程"""
        while self.active:
            try:
                dbName, collectionName, d = self.queue.get(block=True, timeout=1)
                try:
                    self.mainEngine.dbInsert(dbName, collectionName, d)
                except DuplicateKeyError:
                    pass
            except Empty:
                pass

    def registerEvent(self):
        """注册事件回调函数"""
        self.eventEngine.register(EVENT_TIMER, self.processTimerEvent)
        self.eventEngine.register(EVENT_OM_VIX, self.processOmVixEvent)

    def processTimerEvent(self, event):
        """处理定时事件"""
        for calculator in self.vixDict.values():
            self.calcVix(calculator)

    def processOmVixEvent(self, event):
        """处理波动率数据事件"""
        vix = event.dict_['data']
        self.onVixTick(vix)

        if self.bg:
            self.bg.updateTick(vix)

        print(vix.vtSymbol, vix.datetime.strftime('%Y%m%d %H:%M:%S.%f'), vix.lastPrice)

    def onVixTick(self, vix):
        """波动率tick更新"""
        self.insertData(TICK_DB_NAME, vix.vtSymbol, vix)

    def onVixBar(self, bar):
        """波动率bar更新"""
        self.insertData(MINUTE_DB_NAME, bar.vtSymbol, bar)

    def insertData(self, dbName, collectionName, data):
        """把数据放入入库队列"""
        self.queue.put((dbName, collectionName, data.__dict__))

    def calcVix(self, calculator):
        """计算波指"""
        if self.active:
            tName = 'vixT.' + calculator.chain.symbol
            fIdxName = 'vixFIdx.' + calculator.chain.symbol
            vixName = 'vix.' + calculator.chain.symbol

            option = calculator.chain.optionDict.values()[0]
            if not calculator.t:
                self.timer(tName, 1, calculator.calcT, option)
            else:
                self.timer(tName, 60, calculator.calcT, option)

            if not calculator.fIdx:
                self.timer(fIdxName, 1, calculator.calcFIdx)
            else:
                self.timer(fIdxName, 12, calculator.calcFIdx)

            vix = self.timer(vixName, 0, calculator.calcVix)
            if vix:
                event = Event(type_=EVENT_OM_VIX)
                event.dict_['data'] = vix
                self.eventEngine.put(event)

    def timer(self, name, interval, func, *args, **kwargs):
        """自定义定时器"""
        if self.timerDict.get(name) is None:
            self.timerDict[name] = 0

        if self.timerDict[name] > interval:
            res = func(*args, **kwargs)
            self.timerDict[name] = 0
            return res

        self.timerDict[name] += 1


########################################################################
class OmStrategyEngine(object):
    """策略引擎"""
    settingFileName = 'strategy_setting.json'
    settingFilePath = getJsonPath(settingFileName, __file__)

    # ----------------------------------------------------------------------
    def __init__(self, omEngine, eventEngine):
        """Constructor"""
        self.omEngine = omEngine
        self.mainEngine = omEngine.mainEngine
        self.eventEngine = eventEngine

        self.portfolio = None

        self.today = todayDate()

        self.strategyDict = {}  # name: strategy
        self.symbolStrategyDict = {}  # vtSymbol：strategy list
        self.orderStrategyDict = {}  # vtOrderID: strategy

        self.registerEvent()

        print('Strategy Engine run')

    # ----------------------------------------------------------------------
    def registerEvent(self):
        """注册事件监听"""
        self.eventEngine.register(EVENT_TICK, self.processTickEvent)
        self.eventEngine.register(EVENT_TRADE, self.processTradeEvent)
        self.eventEngine.register(EVENT_ORDER, self.processOrderEvent)
        self.eventEngine.register(EVENT_TIMER, self.processTimerEvent)
        self.eventEngine.register(EVENT_OM_VIX, self.processVixTickEvent)

    # ----------------------------------------------------------------------
    def writeLog(self, content):
        """快速发出日志事件"""
        log = VtLogData()
        log.logContent = content

        event = Event(EVENT_OM_STRATEGYLOG)
        event.dict_['data'] = log
        self.eventEngine.put(event)

    # ----------------------------------------------------------------------
    def callStrategyFunc(self, strategy, func, params=None):
        """调用策略的函数，若触发异常则捕捉"""
        try:
            if params:
                func(params)
            else:
                func()
        except Exception:
            # 停止策略，修改状态为未初始化
            strategy.trading = False
            strategy.inited = False

            # 发出日志
            content = '\n'.join([u'策略%s触发异常已停止' % strategy.name,
                                 traceback.format_exc()])
            self.writeLog(content)

    # ----------------------------------------------------------------------
    def processVixTickEvent(self, event):
        vixTick = event.dict_['data']
        l = self.strategyDict.values()

        for strategy in l:
            if strategy.inited:
                self.callStrategyFunc(strategy, strategy.onVixTick, vixTick)


    # ----------------------------------------------------------------------
    def processTickEvent(self, event):
        """处理行情事件"""
        tick = event.dict_['data']
        l = self.symbolStrategyDict.get(tick.vtSymbol, None)

        if l:
            for strategy in l:
                if strategy.inited:
                    self.callStrategyFunc(strategy, strategy.onTick, tick)

    # ----------------------------------------------------------------------
    def processTradeEvent(self, event):
        """处理成交事件"""
        trade = event.dict_['data']
        strategy = self.orderStrategyDict.get(trade.vtOrderID, None)
        if strategy:
            self.callStrategyFunc(strategy, strategy.onTrade, trade)

    # ----------------------------------------------------------------------
    def processOrderEvent(self, event):
        """处理委托事件"""
        order = event.dict_['data']
        strategy = self.orderStrategyDict.get(order.vtOrderID, None)
        if strategy:
            self.callStrategyFunc(strategy, strategy.onOrder, order)

    # ----------------------------------------------------------------------
    def processTimerEvent(self, event):
        """处理定时事件"""
        l = self.strategyDict.values()

        for strategy in l:
            if strategy.inited:
                self.callStrategyFunc(strategy, strategy.onTimer)

    # ----------------------------------------------------------------------
    def loadSetting(self):
        """加载配置"""
        self.portfolio = self.omEngine.portfolio

        with open(self.settingFilePath) as f:
            l = json.load(f)

            for setting in l:
                self.loadStrategy(setting)

    # ----------------------------------------------------------------------
    def loadStrategy(self, setting):
        """加载策略"""
        try:
            name = setting['name']
            className = setting['className']
        except Exception:
            msg = traceback.format_exc()
            self.writeLog(u'载入策略出错：%s' % msg)
            return

        # 获取策略类
        strategyClass = STRATEGY_CLASS.get(className, None)
        if not strategyClass:
            self.writeLog(u'找不到策略类：%s' % className)
            return

        # 防止策略重名
        if name in self.strategyDict:
            self.writeLog(u'策略实例重名：%s' % name)
        else:
            # 创建策略实例
            strategy = strategyClass(self, setting)
            self.strategyDict[name] = strategy

            # 保存Tick映射关系
            for vtSymbol in strategy.vtSymbols:
                print(vtSymbol)
                if vtSymbol in self.symbolStrategyDict:
                    l = self.symbolStrategyDict[vtSymbol]
                else:
                    l = []
                    self.symbolStrategyDict[vtSymbol] = l
                l.append(strategy)
                print(self.symbolStrategyDict)

    # ----------------------------------------------------------------------
    def initStrategy(self, name):
        """初始化策略"""
        strategy = self.strategyDict[name]

        if not strategy.inited:
            strategy.inited = True
            self.callStrategyFunc(strategy, strategy.onInit)

    # ----------------------------------------------------------------------
    def startStrategy(self, name):
        """启动策略"""
        strategy = self.strategyDict[name]

        if strategy.inited and not strategy.trading:
            strategy.trading = True
            self.callStrategyFunc(strategy, strategy.onStart)

    # ----------------------------------------------------------------------
    def stopStrategy(self, name):
        """停止策略"""
        strategy = self.strategyDict[name]

        if strategy.trading:
            strategy.trading = False
            self.callStrategyFunc(strategy, strategy.onStop)

            for vtOrderID, s in self.orderStrategyDict.items():
                if s is strategy:
                    self.cancelOrder(vtOrderID)

    # ----------------------------------------------------------------------
    def loadVixBar(self, dbName, collectionName, days):
        """从数据库中读取Bar数据，startDate是datetime对象"""

        # 如果没有则从数据库中读取数据
        startDate = self.today - timedelta(days)

        d = {'datetime': {'$gte': startDate}}
        barData = self.mainEngine.dbQuery(dbName, collectionName, d, 'datetime')

        l = []
        for d in barData:
            bar = VtBarData()
            bar.__dict__ = d
            l.append(bar)
        return l

    # ----------------------------------------------------------------------
    def loadVixTick(self, dbName, collectionName, days):
        """从数据库中读取Tick数据，startDate是datetime对象"""
        startDate = self.today - timedelta(days)

        d = {'datetime': {'$gte': startDate}}
        tickData = self.mainEngine.dbQuery(dbName, collectionName, d, 'datetime')

        l = []
        for d in tickData:
            tick = VtTickData()
            tick.__dict__ = d
            l.append(tick)
        return l


    # ----------------------------------------------------------------------
    def sendOrder(self, vtSymbol, direction, offset, price, volume):
        """发单"""
        contract = self.mainEngine.getContract(vtSymbol)
        if not contract:
            return ''

        req = VtOrderReq()
        req.symbol = contract.symbol
        req.exchange = contract.exchange
        req.vtSymbol = vtSymbol
        req.price = price
        req.volume = volume
        req.direction = direction
        req.offset = offset
        req.priceType = PRICETYPE_LIMITPRICE

        return self.mainEngine.sendOrder(req, contract.gatewayName)

    # ----------------------------------------------------------------------
    def cancelOrder(self, vtOrderID):
        """撤单"""
        order = self.mainEngine.getOrder(vtOrderID)
        if not order:
            return

        req = VtCancelOrderReq()
        req.symbol = order.symbol
        req.exchange = order.exchange
        req.orderID = order.orderID
        req.symbol = order.symbol
        req.vtSymbol = order.vtSymbo

        self.mainEngine.cancelOrder(req, order.gatewayName)

    # ----------------------------------------------------------------------
    def buy(self, vtSymbol, price, volume):
        """开多"""
        return self.sendOrder(vtSymbol, DIRECTION_LONG, OFFSET_OPEN, price, volume)

    # ----------------------------------------------------------------------
    def short(self, vtSymbol, price, volume):
        """开空"""
        return self.sendOrder(vtSymbol, DIRECTION_SHORT, OFFSET_OPEN, price, volume)

    # ----------------------------------------------------------------------
    def sell(self, vtSymbol, price, volume):
        """平多"""
        return self.sendOrder(vtSymbol, DIRECTION_SHORT, OFFSET_CLOSE, price, volume)

    # ----------------------------------------------------------------------
    def cover(self, vtSymbol, price, volume):
        """平空"""
        return self.sendOrder(vtSymbol, DIRECTION_LONG, OFFSET_CLOSE, price, volume)

    # ----------------------------------------------------------------------
    def dbQuery(self, collectionName, flt):
        """查询数据"""
        return self.mainEngine.dbQuery(OM_DB_NAME, collectionName, flt)

    # ----------------------------------------------------------------------
    def dbUpdate(self, collectionName, d, flt):
        """更新数据"""
        self.mainEngine.dbUpdate(OM_DB_NAME, collectionName, d, flt, True)

    # ----------------------------------------------------------------------
    def getOption(self, vtSymbol):
        """获取期权信息"""
        return self.portfolio.optionDict.get(vtSymbol, None)

    # ----------------------------------------------------------------------
    def getUnderlying(self, vtSymbol):
        """获取标的信息"""
        return self.portfolio.underlyingDict.get(vtSymbol, None)

    # ----------------------------------------------------------------------
    def getChain(self, symbol):
        """获取期权链信息"""
        return self.portfolio.chainDict.get(symbol, None)

    # ----------------------------------------------------------------------
    def getPortfolio(self):
        """获取持仓组合信息"""
        return self.portfolio

    # ----------------------------------------------------------------------
    def putStrategyEvent(self, name):
        """触发策略状态变化事件（通常用于通知GUI更新）"""
        event = Event(EVENT_OM_STRATEGY + name)
        self.eventEngine.put(event)

    # ----------------------------------------------------------------------
    def setStrategyParam(self, name, key, value):
        """设置策略变量"""
        if name in self.strategyDict:
            strategy = self.strategyDict[name]
            strategy.__setattr__(key, value)
            self.writeLog(u'策略%s参数%s已修改为%s' % (name, key, value))
        else:
            self.writeLog(u'策略实例不存在：' + name)
            return None

            # ----------------------------------------------------------------------

    def getStrategyVar(self, name):
        """获取策略当前的变量字典"""
        if name in self.strategyDict:
            strategy = self.strategyDict[name]
            varDict = OrderedDict()

            for key in strategy.varList:
                varDict[key] = strategy.__getattribute__(key)

            return varDict
        else:
            self.writeLog(u'策略实例不存在：' + name)
            return None

    # ----------------------------------------------------------------------
    def getStrategyParam(self, name):
        """获取策略的参数字典"""
        if name in self.strategyDict:
            strategy = self.strategyDict[name]
            paramDict = OrderedDict()

            for key in strategy.paramList:
                paramDict[key] = strategy.__getattribute__(key)

            return paramDict
        else:
            self.writeLog(u'策略实例不存在：' + name)
            return None

            # ----------------------------------------------------------------------

    def initAll(self):
        """全部初始化"""
        for name in self.strategyDict.keys():
            self.initStrategy(name)

    # ----------------------------------------------------------------------
    def startAll(self):
        """全部启动"""
        for name in self.strategyDict.keys():
            self.startStrategy(name)

            # ----------------------------------------------------------------------

    def stopAll(self):
        """全部停止"""
        for name in self.strategyDict.keys():
            self.stopStrategy(name)
